#include <iostream>
#include <cuda_runtime.h>
#include "../common/common.h"
#define DEVICENUM 0
#define SIZE (1 << 10)

struct aosStruct
{
    float a;
    float b;
};

struct soaStruct
{
    float a[SIZE];
    float b[SIZE];
};

__host__ void summaryCPU(float *array_a, float *array_b, float *array_res, int size)
{
    for (int i = 0; i < size; i++)
        array_res[i] = array_a[i] + array_b[i];
}

__global__ void summaryAOSGPU(float *array_a, float *array_b, struct aosStruct *array_res)
{
    int index = threadIdx.x + blockDim.x * blockIdx.x;
    array_res[index].a = array_a[index] + array_b[index];
}

__global__ void summarySOAGPU(float *array_a, float *array_b, struct soaStruct *array_res)
{
    int index = threadIdx.x + blockDim.x * blockIdx.x;
    (array_res->a)[index] = array_a[index] + array_b[index];
}

void checkResult_aos(float *hostRef, struct aosStruct *gpuRef, const int N)
{
    double epsilon = 1.0E-8;
    for (int i = 0; i < N; i++)
    {
        if (abs(hostRef[i] - gpuRef[i].a) > epsilon)
        {
            std::cout << "Results don\'t match!" << std::endl;
            std::cout << hostRef[i] << "(cpu) != " << gpuRef[i].a << "(gpu)"
                      << "|index=" << i << std::endl;
            return;
        }
    }
    std::cout << "Check aos cpu & gpu result success!" << std::endl;
}

void checkResult_soa(float *hostRef, struct soaStruct *gpuRef, const int N)
{
    double epsilon = 1.0E-8;
    for (int i = 0; i < N; i++)
    {
        if (abs(hostRef[i] - (gpuRef->a)[i]) > epsilon)
        {
            std::cout << "Results don\'t match!" << std::endl;
            std::cout << hostRef[i] << "(cpu) != " << (gpuRef->a)[i] << "(gpu)"
                      << "|index=" << i << std::endl;
            return;
        }
    }
    std::cout << "Check soa cpu & gpu result success!" << std::endl;
}

int main(int argc, char **argv)
{
    cudaSetDevice(DEVICENUM);
    int nElem = SIZE;
    size_t nBytes = sizeof(float) * nElem;
    size_t nBytes_aos = sizeof(struct aosStruct) * nElem;
    size_t nBytes_soa = sizeof(struct soaStruct) * nElem;

    double time_cpu, time_aos, time_soa;

    // cpu申请内存并初始化数据
    float *host_array_a = (float *)malloc(nBytes);
    float *host_array_b = (float *)malloc(nBytes);
    float *host_array_res = (float *)malloc(nBytes);
    initialDataRandom(host_array_a, nElem);
    initialDataRandom(host_array_b, nElem);
    memset(host_array_res, 0, nBytes);
    // 计算cpu中两数组相加结果并计算时间
    time_cpu = get_time();
    summaryCPU(host_array_a, host_array_b, host_array_res, nElem);
    CHECK(cudaDeviceSynchronize());
    std::cout << "CPU Sum Array time:" << get_time() - time_cpu << "ms" << std::endl;

    // gpu内存申请
    float *device_array_a = nullptr;
    float *device_array_b = nullptr;
    CHECK(cudaMalloc((float **)&device_array_a, nBytes));
    CHECK(cudaMalloc((float **)&device_array_b, nBytes));
    // 给gpu内存初始化数据
    CHECK(cudaMemcpy(device_array_a, host_array_a, nBytes, cudaMemcpyHostToDevice));
    CHECK(cudaMemcpy(device_array_b, host_array_b, nBytes, cudaMemcpyHostToDevice));

    // 给aos结构的结果数组申请内存并初始化
    struct aosStruct *device_aos_res = nullptr;
    CHECK(cudaMalloc((struct aosStruct **)&device_aos_res, nBytes_aos));
    CHECK(cudaMemset(device_aos_res, 0, nBytes_aos));

    // 给soa结构的结果数组申请内存并初始化
    struct soaStruct *device_soa_res = nullptr;
    CHECK(cudaMalloc((struct soaStruct **)&device_soa_res, nBytes_soa));
    CHECK(cudaMemset(device_soa_res, 0, nBytes_soa));

    // 调用kernel函数执行gpu数组加法运算
    dim3 block_size = block_dims(nElem);
    dim3 grid_size = grid_dims(nElem);
    // cuda warm up
    // for (size_t i = 0; i < 10; i++)
    //     summaryAOSGPU<<<grid_size, block_size>>>(device_array_a, device_array_b, device_aos_res);

    time_aos = get_time();
    summaryAOSGPU<<<grid_size, block_size>>>(device_array_a, device_array_b, device_aos_res);
    CHECK(cudaDeviceSynchronize());
    std::cout << "GPU Sum Array time aos:" << get_time() - time_aos << "ms" << std::endl;

    time_soa = get_time();
    summarySOAGPU<<<grid_size, block_size>>>(device_array_a, device_array_b, device_soa_res);
    CHECK(cudaDeviceSynchronize());
    std::cout << "GPU Sum Array time soa:" << get_time() - time_soa << "ms" << std::endl;

    // 将aos模式下gpu运算结果复制到cpu后并与cpu结果进行比较
    struct aosStruct *res_aos_to_cpu = (struct aosStruct *)malloc(nBytes_aos);
    memset(res_aos_to_cpu, 0, nBytes_aos);
    CHECK(cudaMemcpy(res_aos_to_cpu, device_aos_res, nBytes_aos, cudaMemcpyDeviceToHost));
    checkResult_aos(host_array_res, res_aos_to_cpu, nElem);

    // 将soa模式下gpu运算结果复制到cpu后并与cpu结果进行比较
    struct soaStruct *res_soa_to_cpu = (struct soaStruct *)malloc(nBytes_soa);
    memset(res_soa_to_cpu, 0, nBytes_soa);
    CHECK(cudaMemcpy(res_soa_to_cpu, device_soa_res, nBytes_soa, cudaMemcpyDeviceToHost));
    checkResult_soa(host_array_res, res_soa_to_cpu, nElem);

    // 释放gpu和cpu内存
    cudaFree(device_array_a);
    cudaFree(device_array_b);
    cudaFree(device_aos_res);
    cudaFree(device_soa_res);

    free(host_array_a);
    free(host_array_b);
    free(host_array_res);
    free(res_aos_to_cpu);
    free(res_soa_to_cpu);

    return 0;
}